import logging
from pathlib import Path, PurePath
from typing import Optional

import pypsrp
import spnego
from pypsrp.client import Client
from pypsrp.exceptions import AuthenticationError  # noqa: F401
from pypsrp.powershell import PowerShell, RunspacePool
from urllib3 import connectionpool

from common.credentials import Credentials, LMHash, NTHash, Password, Username, get_plaintext
from infection_monkey.i_puppet import TargetHost

from .powershell_authentication_options import AuthenticationOptions

logger = logging.getLogger(__name__)


def _set_sensitive_packages_log_level_to_error():
    # If root logger is inherited, extensive and potentially sensitive info could be logged
    sensitive_packages = [pypsrp, spnego, connectionpool]
    for package in sensitive_packages:
        logging.getLogger(package.__name__).setLevel(logging.ERROR)


# The pypsrp library requires LM or NT hashes to be formatted like "LM_HASH:NT_HASH"
#
# Example:
# If your LM hash is 1ec78eb5f6edd379351858c437fc3e4e and your NT hash is
# 79a760336ad8c808fee32aa96985a305, then you would pass
# "1ec78eb5f6edd379351858c437fc3e4e:79a760336ad8c808fee32aa96985a305" as the
# `password` parameter to pypsrp.
#
# In our case, we have a set of NT hashes and a set of LM hashes, but we don't
# know if any particular LM/NT hash pair was generated from the same password.
# To avoid confusion, we pair each NT or LM hash with a dummy (i.e. all zeros)
# hash.
def format_password(credentials: Credentials) -> Optional[str]:
    secret = credentials.secret

    if not secret:
        return secret

    if isinstance(secret, Password):
        plaintext_secret = get_plaintext(secret.password)
        return plaintext_secret

    if isinstance(secret, LMHash):
        plaintext_secret = get_plaintext(secret.lm_hash)
        return f"{plaintext_secret}:00000000000000000000000000000000"

    if isinstance(secret, NTHash):
        plaintext_secret = get_plaintext(secret.nt_hash)
        return f"00000000000000000000000000000000:{plaintext_secret}"

    raise ValueError(f"Unknown secret type {type(secret)}")


class PowerShellClient:
    """
    A client for executing commands on a remote host using PowerShell.
    """

    def __init__(self):
        _set_sensitive_packages_log_level_to_error()

        self._client = None
        self._authenticated = False

    def connect(
        self,
        host: TargetHost,
        credentials: Credentials,
        auth_options: AuthenticationOptions,
        timeout: float,
    ):
        """
        Connects to the remote host using the given credentials.

        :param host: The host to connect to
        :param credentials: The credentials to use
        :param auth_options: The authentication options to use
        :param timeout: The timeout for the connection
        :raises Exception: If an error occurred while attempting to connect
        """
        username = (
            credentials.identity.username
            if isinstance(credentials.identity, Username)
            else credentials.identity
        )
        self._client = Client(
            str(host.ip),
            username=username,
            password=format_password(credentials),
            cert_validation=False,
            auth=auth_options.authentication_type.value,
            encryption=auth_options.encryption_setting.value,
            ssl=auth_options.ssl_enabled,
            connection_timeout=timeout,
        )

        # Attempt to execute dir command to know if authentication was successful.
        # This will raise an exception if authentication was not successful.
        self._client.execute_cmd("dir")
        self._authenticated = True
        logger.debug("Successfully authenticated to remote PowerShell service")

    def connected(self) -> bool:
        return self._authenticated

    def copy_file(self, src: Path, dest: PurePath):
        """
        Copies a file from the local machine to the remote machine.

        :param src: The path to the file to copy
        :param dest: The destination path on the remote machine
        :raises Exception: If an error occurred while attempting to copy the file
        """
        try:
            self._client.copy(str(src), str(dest))  # type: ignore[union-attr]
            logger.debug(f"Successfully copied {src}")
        except Exception as err:
            logger.error(f"Failed to copy {src} to {dest}: {err}")
            raise err

    def execute_cmd_as_detached_process(self, cmd: str):
        """
        Executes a command on the remote host. The command will be executed in detached process

        :param cmd: The command to execute
        :raises Exception: If an error occurred while attempting to execute the command
        """
        logger.debug("Attempting to execute a command on the remote host as a detached process.")
        with (
            self._client.wsman,  # type: ignore[union-attr]
            RunspacePool(self._client.wsman) as pool,  # type: ignore[union-attr]
        ):
            ps = PowerShell(pool)
            ps.add_cmdlet("Invoke-WmiMethod").add_parameter("path", "win32_process").add_parameter(
                "name", "create"
            ).add_parameter("ArgumentList", cmd)
            ps.invoke()
